import argparse
import torch
import os
import numpy as np
from hierarchical_runner import HierarchicalRunner
from environment.hierarchical_encirclement_env import HierarchicalEncirclementEnv
from common.arguments import get_args

def main():
    # 获取参数
    args = get_args()
    
    # 添加分层决策特定参数
    args.num_targets = 3  # 目标数量
    args.reassignment_interval = 10  # 重新分配间隔
    args.scenario_name = 'hierarchical_encirclement'
    
    # 创建分层环境
    env = HierarchicalEncirclementEnv(
        num_uavs=args.num_uavs,
        grid_size=args.grid_size,
        num_position_action=args.num_position_action,
        num_energy_action=args.num_energy_action,
        num_targets=args.num_targets,
        seed=args.seed
    )
    
    # 创建分层训练器
    runner = HierarchicalRunner(args, env)
    
    if args.evaluate:
        # 评估模式
        print("开始评估分层决策模型...")
        # 这里可以添加评估逻辑
    else:
        # 训练模式
        print("开始分层决策训练...")
        rewards, success_rates = runner.hierarchical_train()
        
        print(f"训练完成!")
        print(f"平均奖励: {np.mean(rewards):.2f}")
        print(f"平均成功率: {np.mean(success_rates):.2f}")
        print(f"最终成功率: {success_rates[-1]:.2f}")

if __name__ == '__main__':
    main()